#include <thrust/host_vector.h>
#include <thrust/count.h>
#include <thrust/sort.h>

#include "pt2curve_projection.h"
#include <Algorithm/LBVH/lbvh_cuda.cuh>
#include "lbvh_specs.cuh"
#include "pt2curve_kernels_cuda.cuh"

namespace shermit {
namespace fitting {
namespace curve {
	template<typename Float_t>
	Pt2CurveProjection<Float_t, DeviceType::GPU>::Pt2CurveProjection(
		const std::vector<nv::vec2<Float_t>>& h_pts, Float_t pt_half_length
	)  :pts(h_pts.begin(),h_pts.end()),
		bin_count(h_pts.size(),0),
		tl(h_pts.size()),tr(h_pts.size()),tk(h_pts.size()),found(h_pts.size(),false),
		pts_projected(h_pts.size()),
		projection_dist2(h_pts.size()),
		num_samples(0),tangents(nullptr),normals(nullptr),
		tsamples(nullptr),p_tsamples(nullptr),isects(nullptr),
		use_lbvh(false),lbvh(nullptr)
	{
		if (pt_half_length > 0) {
			use_lbvh = true;
			lbvh = new lbvh::LBVH_CUDA < AABB2_SOA<Float_t>>();
			auto lbvh_handle = (lbvh::LBVH_CUDA < AABB2_SOA<Float_t>>*)lbvh;
			auto bm = BuildManager_Pts<Float_t>();
			bm.half_length = pt_half_length;
			bm.pts = getRawPtr(pts);
			bm.size = pts.size();
			lbvh_handle->init(bm);
		}
	}

	template<typename Float_t>
	Pt2CurveProjection<Float_t, DeviceType::GPU>::~Pt2CurveProjection() {
		deallocate();
		if (lbvh) delete lbvh;
	}

	template<typename Float_t>
	void Pt2CurveProjection<Float_t, DeviceType::GPU>::resizeVecs(int size) {
		if (size > num_samples) {
			deallocate();
			allocate(size);
		}
	}

	template<typename Float_t>
	const nv::vec2<Float_t>* Pt2CurveProjection<Float_t, DeviceType::GPU>::getDataPoints() {
		const nv::vec2<Float_t>* ptr = getRawPtr(pts);
		return ptr;
	}

	template<typename Float_t>
	void Pt2CurveProjection<Float_t, DeviceType::GPU>::allocate(int size) {
		num_samples = size;
		checkCudaErrors(GPUMalloc(tangents, sizeof(nv::vec2<Float_t>)*num_samples));
		checkCudaErrors(GPUMalloc(normals, sizeof(nv::vec2<Float_t>) * num_samples));
		checkCudaErrors(GPUMalloc(tsamples, sizeof(Float_t) * num_samples));
		checkCudaErrors(GPUMalloc(p_tsamples, sizeof(nv::vec2<Float_t>) * num_samples));
		checkCudaErrors(GPUMalloc(isects, sizeof(int) * num_samples));
	}

	template<typename Float_t>
	void Pt2CurveProjection<Float_t, DeviceType::GPU>::deallocate() {
		if (num_samples == 0) return;
		num_samples = 0;
		checkCudaErrors(cudaFree(tangents));
		checkCudaErrors(cudaFree(normals));
		checkCudaErrors(cudaFree(tsamples));
		checkCudaErrors(cudaFree(p_tsamples));
		checkCudaErrors(cudaFree(isects));
	}

	template<typename Float_t>
	Float_t Pt2CurveProjection<Float_t, DeviceType::GPU>::computeProjectionDirect(
		const CubicBSplineCurve<Float_t, DeviceType::GPU>& curve,
		int max_iter, int spr) 
	{
		assert(max_iter > 0);
		assert(spr <= 1024);
		assert(spr > 0);
		
		int num_pts = pts.size();
		int num_samples = num_pts * spr;
		resizeVecs(num_samples);
		
		// re init
		//thrust::fill(bin_count.begin(), bin_count.end(), 0);
		thrust::fill(found.begin(), found.end(), false);
		
		// generate initial curve lines
		thrust::fill(tl.begin(), tl.end(), curve.tMin());
		thrust::fill(tr.begin(), tr.end(), curve.tMax());

		int num_not_converged = thrust::count(found.begin(), found.end(), false);
		int _iter = 0;
		while (num_not_converged > 0 && _iter < max_iter) {
			if (0) {
				thrust::host_vector<bool> h_found = found;
				thrust::host_vector<Float_t> h_tk = tk;
				thrust::host_vector<Float_t> h_tl = tl;
				thrust::host_vector<Float_t> h_tr = tr;
				printf("iter: %d\n", _iter);
				for (int i = 0; i < num_pts; ++i) {
					printf("found: %d, t: %lf range: (%lf, %lf)\n", h_found[i], h_tk[i], h_tl[i], h_tr[i]);
				}
			}
			sampleCubicBSplineCurveLinesLocal << <num_pts, spr >> > (
				curve,
				spr, getRawPtr(found), getRawPtr(tl), getRawPtr(tr),
				tsamples, p_tsamples,
				tangents, normals
				);
			cudaDeviceSynchronize();
			getLastCudaError("sampleCubicBSplineCurveLinesLocal kernel failure");

			pickClosestTksDirect << <num_pts, spr, sizeof(Float_t)* spr * 4 >> > (
				curve,
				getRawPtr(pts),
				tsamples, p_tsamples,
				tangents, normals,
				getRawPtr(tk), getRawPtr(pts_projected),
				getRawPtr(found), getRawPtr(tl), getRawPtr(tr)
				);
			cudaDeviceSynchronize();
			getLastCudaError("pickClosestTksDirect kernel failure");

			num_not_converged = thrust::count(found.begin(), found.end(), false);
			_iter++;
		}

		// calculate root mean square error
		auto& CudaDevice = MCudaDevice::getInstance();
		CudaDevice.launchKernel(num_pts, 0,
			calcProjectionDistance2,
			num_pts, getRawPtr(pts), getRawPtr(pts_projected), getRawPtr(projection_dist2));
		Float_t rmse = thrust::reduce(projection_dist2.begin(), projection_dist2.end(), 0.0);
		rmse = sqrt(rmse / num_pts);

		// report results
		if (0) {
			thrust::host_vector<bool> h_found = found;
			thrust::host_vector<Float_t> h_tk = tk;
			thrust::host_vector<nv::vec2<Float_t>> h_pts = pts;
			printf("iter: %d\n", _iter);
			for (int i = 0; i < num_pts; ++i) {
				//if (h_found[i]) continue;
				//printf("(%lf, %lf)  t: %lf\n", h_pts[i].x, h_pts[i].y, h_tk[i]);
				printf("t: %lf\n", h_tk[i]);
			}
		}
		printf("projection finished in %d iters with %d data points not converged.\n", _iter, num_not_converged);
		printf("Root Mean Square Error: %lf\n", rmse);
		return rmse;
	}

	template<typename Float_t>
	Float_t Pt2CurveProjection<Float_t, DeviceType::GPU>::computeProjectionLBVH(
		const CubicBSplineCurve<Float_t, DeviceType::GPU>& curve,
		int max_iter, int spr)
	{
		assert(max_iter > 0);
		assert(spr <= 1024);
		assert(spr > 0);

		int num_pts = pts.size();
		int num_samples = num_pts * spr;
		resizeVecs(num_samples);
		auto lbvh_handle = (lbvh::LBVH_CUDA < AABB2_SOA<Float_t>>*)lbvh;

		// re init
		thrust::fill(bin_count.begin(), bin_count.end(), 0);
		thrust::fill(found.begin(), found.end(), false);

		// generate initial curve lines
		thrust::fill(tl.begin(), tl.end(), curve.tMin());
		thrust::fill(tr.begin(), tr.end(), curve.tMax());
		sampleCubicBSplineCurveLinesLocal << <num_pts, spr >> > (
			curve,
			spr, getRawPtr(found), getRawPtr(tl), getRawPtr(tr),
			tsamples, p_tsamples,
			tangents, normals
			);
		cudaDeviceSynchronize();
		getLastCudaError("sampleCubicBSplineCurveLinesLocal kernel failure");
		/*
		Float_t tl0 = curve.tMin(), tr0 = curve.tMax();
		Float_t step = (tr0 - tl0) / num_samples;
		sampleCubicBSplineCurveGlobal << <num_pts, spr >> > (
			curve, step, tl0, tr0,
			tsamples, p_tsamples, tangents, normals
			);
		cudaDeviceSynchronize();
		getLastCudaError("sampleCubicBSplineCurveGlobal kernel failure");
		*/
		auto qm = QueryManagerDevice_Lines<Float_t>();
		qm.lines.o = p_tsamples;
		qm.lines.d = normals;
		qm.size = num_samples;
		qm.bin_size = spr;
		qm.bin_count = getRawPtr(bin_count);
		qm.isects = isects;
		lbvh_handle->testIntersection(qm);

		pickClosestTksLBVH << <num_pts, spr, sizeof(Float_t)* spr * 4 >> > (
			curve,
			isects, getRawPtr(bin_count),
			getRawPtr(pts),
			tsamples, p_tsamples,
			tangents, normals,
			getRawPtr(tk), getRawPtr(pts_projected),
			getRawPtr(found), getRawPtr(tl), getRawPtr(tr)
			);
		cudaDeviceSynchronize();
		getLastCudaError("pickClosestTksLBVH kernel failure");
		
		int num_not_converged = thrust::count(found.begin(), found.end(), false);
		int _iter = 1;
		while (num_not_converged > 0 && _iter < max_iter) {
			if (0) {
				thrust::host_vector<bool> h_found = found;
				thrust::host_vector<Float_t> h_tk = tk;
				thrust::host_vector<Float_t> h_tl = tl;
				thrust::host_vector<Float_t> h_tr = tr;
				printf("iter: %d\n", _iter);
				for (int i = 0; i < num_pts; ++i) {
					printf("found: %d, t: %lf range: (%lf, %lf)\n", h_found[i], h_tk[i], h_tl[i], h_tr[i]);
				}
			}
			sampleCubicBSplineCurveLinesLocal << <num_pts, spr >> > (
				curve,
				spr, getRawPtr(found), getRawPtr(tl), getRawPtr(tr),
				tsamples, p_tsamples,
				tangents, normals
				);
			cudaDeviceSynchronize();
			getLastCudaError("sampleCubicBSplineCurveLinesLocal kernel failure");

			thrust::fill(bin_count.begin(), bin_count.end(), 0);
			lbvh_handle->testIntersection(qm);

			pickClosestTksLBVH << <num_pts, spr, sizeof(Float_t)* spr * 4 >> > (
				curve,
				isects, getRawPtr(bin_count),
				getRawPtr(pts),
				tsamples, p_tsamples,
				tangents, normals,
				getRawPtr(tk), getRawPtr(pts_projected),
				getRawPtr(found), getRawPtr(tl), getRawPtr(tr)
				);
			cudaDeviceSynchronize();
			getLastCudaError("pickClosestTksLBVH kernel failure");

			num_not_converged = thrust::count(found.begin(), found.end(), false);
			_iter++;
		}
		/* reorder tks and data pts */
		//thrust::sort_by_key(tk.begin(), tk.end(), pts.begin());

		// calculate root mean square error
		auto& CudaDevice = MCudaDevice::getInstance();
		CudaDevice.launchKernel(num_pts, 0, 
			calcProjectionDistance2, 
			num_pts, getRawPtr(pts), getRawPtr(pts_projected), getRawPtr(projection_dist2));
		Float_t rmse = thrust::reduce(projection_dist2.begin(), projection_dist2.end(), 0.0);
		rmse = sqrt(rmse / num_pts);

		// report results
		if (0) {
			thrust::host_vector<bool> h_found = found;
			thrust::host_vector<Float_t> h_tk = tk;
			thrust::host_vector<nv::vec2<Float_t>> h_pts = pts;
			printf("projection iter: %d\n", _iter);
			for (int i = 0; i < num_pts; ++i) {
				//if (h_found[i]) continue;
				//printf("(%lf, %lf)  t: %lf\n", h_pts[i].x, h_pts[i].y, h_tk[i]);
				printf("t: %lf\n", h_tk[i]);
			}
		}
		printf("projection finished in %d iters with %d data points not converged.\n", _iter, num_not_converged);
		printf("Root Mean Square Error: %lf\n", rmse);
		return rmse;
	}

	template<typename Float_t>
	void Pt2CurveProjection<Float_t, DeviceType::GPU>::computeProjection(
		const CubicBSplineCurve<Float_t, DeviceType::GPU>& curve,
		int max_iter, int spr,
		Float_t* rmse, Float_t** _tk, nv::vec2<Float_t>** _pts_projected,
		nv::vec2<Float_t>** _tangents, nv::vec2<Float_t>** _normals) 
	{
		if (use_lbvh) {
			if (rmse != nullptr) *rmse = computeProjectionLBVH(curve, max_iter, spr);
			if (_tk != nullptr) *_tk = getRawPtr(tk);
			if (_pts_projected != nullptr) *_pts_projected = getRawPtr(pts_projected);
			if (_tangents != nullptr) *_tangents = tangents;
			if (_normals != nullptr) *_normals = normals;
			return;
		}
		if (rmse != nullptr) *rmse = computeProjectionDirect(curve, max_iter, spr);
		if (_tk != nullptr) *_tk = getRawPtr(tk);
		if (_pts_projected != nullptr) *_pts_projected = getRawPtr(pts_projected);
		if (_tangents != nullptr) *_tangents = tangents;
		if (_normals != nullptr) *_normals = normals;
	}

	//template class Pt2CurveProjection<float, DeviceType::GPU>;
	template class Pt2CurveProjection<double, DeviceType::GPU>;
}
}
}